import os

def replace_labels(txt_folder, replacements):
    # 获取txt文件夹中的所有文件
    txt_files = [f for f in os.listdir(txt_folder) if f.endswith('.txt')]

    # 遍历每个txt文件
    for txt_file in txt_files:
        # 读取txt文件内容
        with open(os.path.join(txt_folder, txt_file), 'r') as f:
            lines = f.readlines()

        # 替换标签
        modified_lines = []
        for line in lines:
            parts = line.split()
            label = parts[0]

            # 替换标签
            if label in replacements:
                label = replacements[label]

        # 将修改后的内容写入txt文件
        with open(os.path.join(txt_folder, txt_file), 'w') as f:
            f.write('/n'.join(modified_lines))

if __name__=="__main__":

    txt_folder = 'D:/datasets/visDrone/VisDrone2019-DET-val/new_label_0.01'  # txt文件夹路径

    # 标签替换映射
    replacements = {
        '0': '0',
        '1': '0',
        '3': '1',
        '5': '3',
        '8': '2',
    }

    # 执行替换
    replace_labels(txt_folder, replacements)